// Copyright (c) 2023 Presto Labs Pte. Ltd.
// Author: chensili

#include "coin2/exchange/base/log/latency_recorder.h"
#include "coin2/exchange/base/log/strategy_logger.h"

#include <gtest/gtest.h>

namespace coin2::exchange::base::strategy_util {

class LatencyRecorderTest : public ::testing::Test {
 protected:
  LatencyRecorderTest() {
    auto strat_request = coin::proto::StrategyRequestProto();
    strat_request.set_strategy_group("dmm");
    strat_request.set_strategy_name("dmm_spot_pilot");
    std::string log_root = "~/data/strat_proto_log";
    std::string hostname = GetLocalHostname();
    strat_logger_.reset(new StrategyLogger(
        strat_request,
        log_root,
        hostname
    ));
  }

  const LatencyRecorder::map_type* GetMmap() {
    return latency_recorder_->GetMmap();
  }

  const LatencyRecorder::list_type* GetMlist() {
    return latency_recorder_->GetMlist();
  }

  void SetTimestamp(int64_t curr_ts) {
    latency_recorder_->SetTimestamp(curr_ts);}

  void ResetLatencyRecorder(int cap) {
    latency_recorder_.reset(new LatencyRecorder(
        4,
        std::bind(
            &StrategyLogger::WriteLatency,
            strat_logger_.get(),
            std::placeholders::_1,
            std::placeholders::_2,
            std::placeholders::_3)
    ));
  }

  std::unique_ptr<StrategyLogger> strat_logger_;
  std::unique_ptr<LatencyRecorder> latency_recorder_;
};

TEST_F(LatencyRecorderTest, End) {
  ResetLatencyRecorder(4);
  latency_recorder_->Begin(coin::proto::LatencyProto::FEED_LATENCY);
  auto latency_context_id = latency_recorder_->GetCurrentContextId();
  std::string ts("20230401T000000");
  boost::posix_time::ptime feed_pub_time(boost::posix_time::from_iso_string(ts));
  auto acct_req = coin::proto::AccountRequestProto();
  acct_req.set_market_type("Spot");
  acct_req.set_exchange("Binance");
  acct_req.set_api_version("v1");
  latency_recorder_->SetProtoAndTime(
      latency_context_id,
      coin::proto::LifeOfSignal::RAW_BOOK_FEED_EXCHANGE_PUBLISHED,
      coin2::exchange::base::symbology::DatetimeToTimestamp(feed_pub_time),
      acct_req);
  EXPECT_EQ(GetMmap()->size(), 1);
  EXPECT_EQ(GetMlist()->size(), 1);
  latency_recorder_->End(latency_context_id);
  EXPECT_EQ(GetMmap()->size(), 0);
  EXPECT_EQ(GetMlist()->size(), 1);
}

TEST_F(LatencyRecorderTest, EvictFeed) {
  const size_t cap = 4;
  ResetLatencyRecorder(cap);
  std::string ts("20230401T000000");
  boost::posix_time::ptime feed_pub_time(boost::posix_time::from_iso_string(ts));
  for (int i = 1; i <= cap; ++i) {
    latency_recorder_->Begin(coin::proto::LatencyProto::FEED_LATENCY);
    auto latency_context_id = latency_recorder_->GetCurrentContextId();
    feed_pub_time += boost::posix_time::seconds(1);
    auto acct_req = coin::proto::AccountRequestProto();
    latency_recorder_->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::RAW_BOOK_FEED_EXCHANGE_PUBLISHED,
        coin2::exchange::base::symbology::DatetimeToTimestamp(feed_pub_time),
        acct_req);
  }
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto* m_map = GetMmap();
  auto* m_list = GetMlist();
  EXPECT_EQ(m_map->size(), cap);
  auto order_context_id = latency_recorder_->GetCurrentContextId();
  EXPECT_EQ(m_map->count(order_context_id), 1);
  EXPECT_EQ(m_list->front(), order_context_id);
}

TEST_F(LatencyRecorderTest, EvictOrder) {
  const size_t cap = 4;
  ResetLatencyRecorder(cap);
  std::string ts("20230401T000000");
  boost::posix_time::ptime curr_time(boost::posix_time::from_iso_string(ts));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto order_context_id = latency_recorder_->GetCurrentContextId();
  auto acct_req = coin::proto::AccountRequestProto();
  latency_recorder_->SetProtoAndTime(
      order_context_id,
      coin::proto::LifeOfSignal::BEFORE_OG_ORDER_SUBMIT,
      coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
      acct_req
  );
  for (int i = 1; i <= cap - 1; ++i) {
    latency_recorder_->Begin(coin::proto::LatencyProto::FEED_LATENCY);
    auto latency_context_id = latency_recorder_->GetCurrentContextId();
    curr_time += boost::posix_time::seconds(1);
    auto acct_req = coin::proto::AccountRequestProto();
    latency_recorder_->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::RAW_BOOK_FEED_EXCHANGE_PUBLISHED,
        coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
        acct_req); 
  }
  SetTimestamp(coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto* m_map = GetMmap();
  EXPECT_EQ(m_map->size(), cap);
  EXPECT_EQ(*GetMlist()->rbegin(), order_context_id);
  EXPECT_EQ(m_map->count(order_context_id), 1);
}

TEST_F(LatencyRecorderTest, EvictExpiredOrder) {
  const size_t cap = 4;
  ResetLatencyRecorder(cap);
  std::string ts("20230401T000000");
  boost::posix_time::ptime curr_time(boost::posix_time::from_iso_string(ts));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto order_context_id = latency_recorder_->GetCurrentContextId();
  auto acct_req = coin::proto::AccountRequestProto();
  latency_recorder_->SetProtoAndTime(
      order_context_id,
      coin::proto::LifeOfSignal::BEFORE_OG_ORDER_SUBMIT,
      coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
      acct_req
  );
  for (int i = 1; i <= cap - 1; ++i) {
    latency_recorder_->Begin(coin::proto::LatencyProto::FEED_LATENCY);
    auto latency_context_id = latency_recorder_->GetCurrentContextId();
    curr_time += boost::posix_time::seconds(1);
    auto acct_req = coin::proto::AccountRequestProto();
    latency_recorder_->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::RAW_BOOK_FEED_EXCHANGE_PUBLISHED,
        coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
        acct_req); 
  }
  curr_time += boost::posix_time::minutes(1);
  SetTimestamp(coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto* m_map = GetMmap();
  EXPECT_EQ(m_map->size(), cap);
  EXPECT_EQ(m_map->count(order_context_id), 0);
}

TEST_F(LatencyRecorderTest, ForceEvictOrder) {
  const size_t cap = 4;
  ResetLatencyRecorder(cap);
  std::string ts("20230401T000000");
  boost::posix_time::ptime curr_time(boost::posix_time::from_iso_string(ts));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto first_order_context_id = latency_recorder_->GetCurrentContextId();
  auto acct_req = coin::proto::AccountRequestProto();
  latency_recorder_->SetProtoAndTime(
        first_order_context_id,
        coin::proto::LifeOfSignal::BEFORE_OG_ORDER_SUBMIT,
        coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
        acct_req
  );
  for (int i = 1; i <= cap - 1; ++i) {
    latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
    auto latency_context_id = latency_recorder_->GetCurrentContextId();
    curr_time += boost::posix_time::seconds(1);
    latency_recorder_->SetProtoAndTime(
        latency_context_id,
        coin::proto::LifeOfSignal::BEFORE_OG_ORDER_SUBMIT,
        coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time),
        acct_req
    );
  }
  SetTimestamp(coin2::exchange::base::symbology::DatetimeToTimestamp(curr_time));
  latency_recorder_->Begin(coin::proto::LatencyProto::ORDER_LATENCY);
  auto* m_map = GetMmap();
  EXPECT_EQ(m_map->count(first_order_context_id), 0);
}

}  // namespace coin2::exchange::base::strategy_util